package com.xs.langchain4j_springboot.config;

import com.xs.langchain4j_springboot.service.ToolsService;
import com.xs.langchain4j_springboot.test.WeatherTools;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;

import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonStringSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiChatRequestParameters;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.ToolExecutor;
import dev.langchain4j.service.tool.ToolProvider;
import dev.langchain4j.service.tool.ToolProviderResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.internal.Json.fromJson;

@Slf4j
@Configuration
public class FootballAiConfig {

    /**
     * 尝试通过lowlevel 调用大模型服务
     * @param args
     */
    public static void main(String[] args) {


        OpenAiChatModel model = OpenAiChatModel.builder()
                .apiKey("26f9ab1f-3509-4f8e-b06e-7495bc93373a")
                .modelName("doubao-1.5-pro-32k-250115")
                .baseUrl("https://ark.cn-beijing.volces.com/api/v3")
                .parallelToolCalls(true)
                .build();


        //自定义选择工具
        ToolSpecification toolSpecification = ToolSpecification.builder()
                .name("getWeather")
                .description("Returns the weather forecast for a given city")
                .parameters(JsonObjectSchema.builder()
                        .addStringProperty("city", "The city for which the weather forecast should be returned")
                        .addEnumProperty("temperatureUnit", List.of("CELSIUS", "FAHRENHEIT"))
                        .required("city") // the required properties should be specified explicitly
                        .build())
                .build();


        List<ToolSpecification> toolSpecifications = Arrays.asList(toolSpecification);


        OpenAiChatRequestParameters parameters = OpenAiChatRequestParameters.builder()
                .modelName("doubao-1.5-pro-32k-250115")
                .toolSpecifications(toolSpecifications)
                .build();


        UserMessage userMessage = UserMessage.from("What will the weather be like in London tomorrow?");

        ChatRequest request = ChatRequest.builder()
                .messages(userMessage)
                .parameters(parameters)
                .build();

        ChatResponse response = model.doChat(request);
        AiMessage aiMessage = response.aiMessage();
        System.out.println(aiMessage.text());

        //List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(WeatherTools.class);
        if (aiMessage.hasToolExecutionRequests()) {

            for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {

                //TODO 需要单独根据ToolExecutionRequest 去调用工具


                String result = "It is expected to rain in London tomorrow.";
                ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(toolExecutionRequest, result);

                ChatRequest request2 = ChatRequest.builder()
                        .messages(List.of(userMessage, aiMessage, toolExecutionResultMessage))
                        .parameters(parameters)
                        .build();
                ChatResponse response2 = model.doChat(request2);
                AiMessage aiMessage2 = response2.aiMessage();
                System.out.println(aiMessage2.text());

            }


        }


    }


    /**
     * 会话隔离、会话记忆、流式输出、functionCall
     *
     * @param streamingChatLanguageModel
     * @param toolsService
     * @return
     */
    @Bean
    public FootBallAssistant footBallAssistant(StreamingChatLanguageModel streamingChatLanguageModel, ToolsService toolsService) {


        PersistentChatMemoryStore store = new PersistentChatMemoryStore();
        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(store)
                .build();


        ToolProvider toolProvider = (toolProviderRequest) -> {
            if (toolProviderRequest.userMessage().singleText().contains("booking")) {
                ToolSpecification toolSpecification = ToolSpecification.builder()
                        .name("get_booking_details")
                        .description("Returns booking details")
                        .parameters(JsonObjectSchema.builder()
                                .addStringProperty("bookingNumber")
                                .build())
                        .build();
                ToolExecutor toolExecutor = (toolExecutionRequest, memoryId) -> {
                    String arguments = toolExecutionRequest.arguments();
                    String name = toolExecutionRequest.name();
                    //根据模型名称、参数 手动执行
                    log.info("=====================" + name + "工具执行成功，参数：" + arguments);

                    return name + "工具执行成功，参数：" + arguments;
                };

                return ToolProviderResult.builder()
                        .add(toolSpecification, toolExecutor)
                        .build();
            } else {
                return null;
            }
        };


        return AiServices.builder(FootBallAssistant.class)
                //.tools(toolsService)
                .toolProvider(toolProvider)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }

}
